#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库操作模块
"""

import sqlite3
import os
from typing import List, Tuple
from datetime import datetime


class Database:
    def __init__(self, db_path: str = "~/.smartledger/finance.db"):
        self.db_path = os.path.expanduser(db_path)
        self.init_db()

    def init_db(self):
        """初始化数据库表"""
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建transactions表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS transactions (
                id INTEGER PRIMARY KEY,
                timestamp DATETIME NOT NULL,
                amount REAL NOT NULL,
                type TEXT CHECK(type IN ('收入', '支出', '转账')),
                category TEXT,
                payee TEXT,
                source TEXT,
                note TEXT
            )
        """)
        
        # 为transactions表创建索引以提高查询性能
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_transactions_timestamp 
            ON transactions (timestamp)
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_transactions_type 
            ON transactions (type)
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_transactions_category 
            ON transactions (category)
        """)
        
        # 创建一级分类表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS categories (
                id INTEGER PRIMARY KEY,
                name TEXT UNIQUE,
                color TEXT
            )
        """)
        
        # 创建二级分类表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS subcategories (
                id INTEGER PRIMARY KEY,
                name TEXT,
                parent_id INTEGER,
                color TEXT,
                FOREIGN KEY(parent_id) REFERENCES categories(id)
            )
        """)
        
        # 创建budgets表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS budgets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                category_id INTEGER,
                subcategory_id INTEGER DEFAULT NULL,
                month TEXT,
                amount REAL,
                FOREIGN KEY(category_id) REFERENCES categories(id),
                FOREIGN KEY(subcategory_id) REFERENCES subcategories(id)
            )
        """)
        
        conn.commit()
        conn.close()
        
    def delete_budget(self, category):
        """删除预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 解析分类名称，检查是否为二级分类
        if '|' in category:
            # 二级分类
            parts = category.split('|')
            primary_category = parts[0]
            secondary_category = parts[1]
            
            # 获取一级分类ID
            cursor.execute("SELECT id FROM categories WHERE name = ?", (primary_category,))
            category_id = cursor.fetchone()
            
            # 获取二级分类ID
            subcategory_id = None
            if category_id:
                category_id = category_id[0]
                cursor.execute("SELECT id FROM subcategories WHERE name = ? AND parent_id = ?", 
                              (secondary_category, category_id))
                subcategory_result = cursor.fetchone()
                if subcategory_result:
                    subcategory_id = subcategory_result[0]
            
            # 删除预算
            if subcategory_id:
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                cursor.execute("""
                    DELETE FROM budgets 
                    WHERE category_id = ? AND subcategory_id = ? AND month = ?
                """, (category_id, subcategory_id, current_month))
                
                conn.commit()
        else:
            # 一级分类
            cursor.execute("SELECT id FROM categories WHERE name = ?", (category,))
            category_id = cursor.fetchone()
            
            if category_id:
                category_id = category_id[0]
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                
                # 删除预算
                cursor.execute("""
                    DELETE FROM budgets 
                    WHERE category_id = ? AND subcategory_id IS NULL AND month = ?
                """, (category_id, current_month))
                
                conn.commit()
        
        conn.close()
        
    def update_budget(self, category, amount):
        """更新预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 解析分类名称，检查是否为二级分类
        if '|' in category:
            # 二级分类
            parts = category.split('|')
            primary_category = parts[0]
            secondary_category = parts[1]
            
            # 获取一级分类ID
            cursor.execute("SELECT id FROM categories WHERE name = ?", (primary_category,))
            category_id = cursor.fetchone()
            
            # 获取二级分类ID
            subcategory_id = None
            if category_id:
                category_id = category_id[0]
                cursor.execute("SELECT id FROM subcategories WHERE name = ? AND parent_id = ?", 
                              (secondary_category, category_id))
                subcategory_result = cursor.fetchone()
                if subcategory_result:
                    subcategory_id = subcategory_result[0]
            
            # 更新预算
            if subcategory_id:
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                cursor.execute("""
                    UPDATE budgets 
                    SET amount = ?
                    WHERE category_id = ? AND subcategory_id = ? AND month = ?
                """, (amount, category_id, subcategory_id, current_month))
                
                conn.commit()
        else:
            # 一级分类
            cursor.execute("SELECT id FROM categories WHERE name = ?", (category,))
            category_id = cursor.fetchone()
            
            if category_id:
                category_id = category_id[0]
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                
                # 更新预算
                cursor.execute("""
                    UPDATE budgets 
                    SET amount = ?
                    WHERE category_id = ? AND subcategory_id IS NULL AND month = ?
                """, (amount, category_id, current_month))
                
                conn.commit()
        
        conn.close()
        
    def add_budget(self, category, amount):
        """添加预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 解析分类名称，检查是否为二级分类
        if '|' in category:
            # 二级分类
            parts = category.split('|')
            primary_category = parts[0]
            secondary_category = parts[1]
            
            # 获取一级分类ID
            cursor.execute("SELECT id FROM categories WHERE name = ?", (primary_category,))
            category_id = cursor.fetchone()
            
            # 获取二级分类ID
            subcategory_id = None
            if category_id:
                category_id = category_id[0]
                cursor.execute("SELECT id FROM subcategories WHERE name = ? AND parent_id = ?", 
                              (secondary_category, category_id))
                subcategory_result = cursor.fetchone()
                if subcategory_result:
                    subcategory_id = subcategory_result[0]
            
            # 插入或更新预算
            if subcategory_id:
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                cursor.execute("""
                    INSERT OR REPLACE INTO budgets (category_id, subcategory_id, month, amount)
                    VALUES (?, ?, ?, ?)
                """, (category_id, subcategory_id, current_month, amount))
                
                conn.commit()
        else:
            # 一级分类
            cursor.execute("SELECT id FROM categories WHERE name = ?", (category,))
            category_id = cursor.fetchone()
            
            if category_id:
                category_id = category_id[0]
                # 获取当前月份
                current_month = datetime.now().strftime('%Y-%m')
                
                # 插入或更新预算
                cursor.execute("""
                    INSERT OR REPLACE INTO budgets (category_id, subcategory_id, month, amount)
                    VALUES (?, NULL, ?, ?)
                """, (category_id, current_month, amount))
                
                conn.commit()
        
        conn.close()
        
    def get_monthly_expense(self, category):
        """获取分类月度支出"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取当前月份
        current_month = datetime.now().strftime('%Y-%m')
        
        # 解析分类名称，检查是否为二级分类
        if '|' in category:
            # 对于二级分类，直接使用category字段匹配
            cursor.execute("""
                SELECT SUM(amount) 
                FROM transactions 
                WHERE type = '支出' AND category = ? AND timestamp LIKE ?
            """, (category, f'{current_month}%'))
        else:
            # 对于一级分类，需要匹配所有以该分类开头的记录（包括其二级分类）
            cursor.execute("""
                SELECT SUM(amount) 
                FROM transactions 
                WHERE type = '支出' AND (category = ? OR category LIKE ?) AND timestamp LIKE ?
            """, (category, f'{category}|%', f'{current_month}%'))
        
        result = cursor.fetchone()[0]
        conn.close()
        return result if result else 0.0
        
    def get_budgets(self):
        """获取所有预算"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            SELECT b.month, c.name as category, s.name as subcategory, b.amount
            FROM budgets b
            JOIN categories c ON b.category_id = c.id
            LEFT JOIN subcategories s ON b.subcategory_id = s.id
        """)
        
        results = cursor.fetchall()
        conn.close()
        
        # 构造返回数据，如果是二级分类则显示为"一级分类|二级分类"格式
        budgets = []
        for row in results:
            if row[2]:  # 存在二级分类
                category_name = f"{row[1]}|{row[2]}"
            else:  # 只有一级分类
                category_name = row[1]
            budgets.append({'month': row[0], 'category': category_name, 'amount': row[3]})
        
        return budgets

    def insert_transaction(self, timestamp: str, amount: float, type: str, 
                          category: str, payee: str, source: str, note: str):
        """插入单条交易记录"""
        # 数据验证和清理
        if not timestamp:
            raise ValueError("交易时间不能为空")
        
        # 确保金额是浮点数
        amount = float(amount)
        
        # 验证交易类型
        if type not in ['收入', '支出', '转账']:
            raise ValueError(f"无效的交易类型: {type}")
        
        # 确保其他字段是字符串类型
        category = str(category) if category else '其他'
        payee = str(payee) if payee else ''
        source = str(source) if source else ''
        note = str(note) if note else ''
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT INTO transactions 
            (timestamp, amount, type, category, payee, source, note)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """, (timestamp, amount, type, category, payee, source, note))
        
        conn.commit()
        conn.close()
    
    def insert_transactions(self, transactions: List[Tuple]):
        """批量插入交易记录"""
        if not transactions:
            return
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 数据验证和清理
        cleaned_transactions = []
        for transaction in transactions:
            timestamp, amount, type, category, payee, source, note = transaction
            
            # 数据验证和清理
            if not timestamp:
                raise ValueError("交易时间不能为空")
            
            # 确保金额是浮点数
            amount = float(amount)
            
            # 验证交易类型
            if type not in ['收入', '支出', '转账']:
                raise ValueError(f"无效的交易类型: {type}")
            
            # 确保其他字段是字符串类型
            category = str(category) if category else '其他'
            payee = str(payee) if payee else ''
            source = str(source) if source else ''
            note = str(note) if note else ''
            
            cleaned_transactions.append((timestamp, amount, type, category, payee, source, note))
        
        cursor.executemany("""
            INSERT INTO transactions 
            (timestamp, amount, type, category, payee, source, note)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """, cleaned_transactions)
        
        conn.commit()
        conn.close()

    def get_transactions(self, start_date: str = None, end_date: str = None, transaction_type: str = None) -> List[Tuple]:
        """获取交易记录"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if start_date and end_date:
            if transaction_type:
                cursor.execute("""
                    SELECT * FROM transactions 
                    WHERE timestamp BETWEEN ? AND ? AND type = ?
                    ORDER BY timestamp
                """ , (start_date, end_date, transaction_type))
            else:
                cursor.execute("""
                    SELECT * FROM transactions 
                    WHERE timestamp BETWEEN ? AND ?
                    ORDER BY timestamp
                """ , (start_date, end_date))
        else:
            if transaction_type:
                cursor.execute("SELECT * FROM transactions WHERE type = ? ORDER BY timestamp", (transaction_type,))
            else:
                cursor.execute("SELECT * FROM transactions ORDER BY timestamp")
            
        results = cursor.fetchall()
        conn.close()
        return results

    def insert_category(self, name: str, color: str):
        """插入一级分类"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT OR IGNORE INTO categories (name, color)
            VALUES (?, ?)
        """, (name, color))
        
        conn.commit()
        conn.close()
        
    def insert_subcategory(self, name: str, parent_id: int, color: str):
        """插入二级分类"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            INSERT OR IGNORE INTO subcategories (name, parent_id, color)
            VALUES (?, ?, ?)
        """, (name, parent_id, color))
        
        conn.commit()
        conn.close()

    def get_categories(self) -> List[Tuple]:
        """获取所有一级分类"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("SELECT * FROM categories")
        results = cursor.fetchall()
        conn.close()
        return results
        
    def get_subcategories(self, parent_id: int = None) -> List[Tuple]:
        """获取二级分类，如果提供了parent_id，则只返回该一级分类下的二级分类"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if parent_id is not None:
            cursor.execute("SELECT * FROM subcategories WHERE parent_id = ?", (parent_id,))
        else:
            cursor.execute("SELECT * FROM subcategories")
            
        results = cursor.fetchall()
        conn.close()
        return results
        
    def get_category_tree(self) -> List[dict]:
        """获取分类树结构"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 获取所有一级分类
        cursor.execute("SELECT id, name, color FROM categories")
        categories = cursor.fetchall()
        
        # 构建分类树
        category_tree = []
        for category in categories:
            category_dict = {
                'id': category[0],
                'name': category[1],
                'color': category[2],
                'subcategories': []
            }
            
            # 获取该一级分类下的所有二级分类
            cursor.execute("SELECT id, name, color FROM subcategories WHERE parent_id = ?", (category[0],))
            subcategories = cursor.fetchall()
            
            for subcategory in subcategories:
                category_dict['subcategories'].append({
                    'id': subcategory[0],
                    'name': subcategory[1],
                    'color': subcategory[2]
                })
                
            category_tree.append(category_dict)
            
        conn.close()
        return category_tree
        
    def export_all_data(self):
        """导出所有数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 导出transactions表
        cursor.execute("SELECT * FROM transactions")
        transactions = cursor.fetchall()
        
        # 导出categories表
        cursor.execute("SELECT * FROM categories")
        categories = cursor.fetchall()
        
        # 导出budgets表
        cursor.execute("SELECT * FROM budgets")
        budgets = cursor.fetchall()
        
        conn.close()
        
        return {
            'transactions': transactions,
            'categories': categories,
            'budgets': budgets
        }
        
    def clear_all_data(self):
        """清除所有数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 清空所有数据
        cursor.execute("DELETE FROM transactions")
        cursor.execute("DELETE FROM categories")
        cursor.execute("DELETE FROM budgets")
        
        conn.commit()
        conn.close()
        
        # 重新初始化数据库表
        self.init_db()

    def import_all_data(self, data):
        """导入所有数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 清空现有数据
        cursor.execute("DELETE FROM transactions")
        cursor.execute("DELETE FROM categories")
        cursor.execute("DELETE FROM budgets")
        
        # 导入transactions数据
        for row in data['transactions']:
            # 跳过时间为空或金额为0的交易
            if not row[1] or row[2] == 0:
                continue
            cursor.execute("""
                INSERT INTO transactions 
                (id, timestamp, amount, type, category, payee, source, note)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, row)
            
        # 导入categories数据
        for row in data['categories']:
            cursor.execute("""
                INSERT INTO categories 
                (id, name, color)
                VALUES (?, ?, ?)
            """, row)
            
        # 导入budgets数据
        for row in data['budgets']:
            cursor.execute("""
                INSERT INTO budgets 
                (category_id, month, amount)
                VALUES (?, ?, ?)
            """, row)
            
        conn.commit()
        conn.close()